package leetcode.weekly.week312;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

//Solution4Test
public class Solution4 {

	int[] vals;
	List<List<Integer>> nexts;

	class Info {
		int ans;
		Map<Integer, Integer> ends;

		public Info() {
			ends = new HashMap<>();
		}
	}

	public int numberOfGoodPaths(int[] vals, int[][] edges) {
		this.vals = vals;
		int n = vals.length;
		nexts = new ArrayList<>();
		for (int i = 0; i < n; i++) {
			nexts.add(new ArrayList<>());
		}
		for (int i = 0; i < edges.length; i++) {
			int a = edges[i][0], b = edges[i][1];
			nexts.get(a).add(b);
			if (!nexts.get(a).contains(b)) {
				nexts.get(a).add(b);
			}
			if (!nexts.get(b).contains(a)) {
				nexts.get(b).add(a);
			}
		}
		Info info = this.f(-1, 0, vals[0]);
		return info.ans + n;
	}

	private Info f(int root, int index, int rootVal) {
		List<Integer> set = nexts.get(index);
		Info ans = new Info();
		int val = vals[index];
		ans.ends.put(val, 1);
		for (int next : set) {
			if (next == root) {
				continue;
			}
			Info N = this.f(index, next, vals[next]);//
			ans.ans += N.ans;
			for (Map.Entry<Integer, Integer> entry : N.ends.entrySet()) {
				int key = entry.getKey(), value = entry.getValue();
				if (entry.getKey() >= rootVal) {
					if (ans.ends.containsKey(key)) {
						ans.ans += value * ans.ends.get(key);
					}
					ans.ends.put(key, ans.ends.getOrDefault(key, 0) + value);
				}
			}
		}
		return ans;
	}

}
